package org.example.sse;

import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

@Slf4j
public class SSEServer {
    private static AtomicInteger clientCount = new AtomicInteger(0);

    private static Map<String, SseEmitter> clients = new ConcurrentHashMap();

    public static SseEmitter connect(String userId) {
        SseEmitter sseEmitter = new SseEmitter();
        sseEmitter.onCompletion(onCompletionCallback(userId));
        sseEmitter.onError(onErrorCallback(userId));
        sseEmitter.onTimeout(onTimeoutHandle(userId));
        log.info("创建新的 SSE 连接，用户：{}", userId);
        clients.put(userId, sseEmitter);
        clientCount.getAndIncrement();
        return sseEmitter;
    }

    public static void stopServer(String userId) {
        if(CollectionUtils.isEmpty(clients)) {
            return;
        }

        SseEmitter sseEmitter = clients.get(userId);
        if(sseEmitter!=null) {
            sseEmitter.complete();
            removeConnection(userId);
            log.info("当前连接已关闭，用户：{}", userId);
        } else {
            log.warn("当前连接无需关闭，请勿重复操作");
        }

    }

    public static void sendMessage(String userId, String message) throws IOException {
        if(CollectionUtils.isEmpty(clients)) {
            return;
        }

        if(clients.containsKey(userId)) {
            SseEmitter sseEmitter = clients.get(userId);
            sendMessageHandler(userId, message, sseEmitter);
        }
    }

    public static void sendMessageToAllUser(String message) {
        if(CollectionUtils.isEmpty(clients)) {
           return;
        }

        clients.forEach((userId, sseEmitter) -> {
            sendMessageHandler(userId, message, sseEmitter);
        });
    }

    private static void sendMessageHandler(String userId, String message, SseEmitter sseEmitter) {
        SseEmitter.SseEventBuilder sendMessage = SseEmitter.event().id(userId).name(userId).data(message);
        try {
            sseEmitter.send(sendMessage);
        } catch (IOException e) {
            log.error("发送消息失败，用户：{}", userId);
            removeConnection(userId);
        }
    }

    private static Runnable onTimeoutHandle(String userId) {
        return () -> {
            log.error("SSE 连接超时，用户：{}", userId);
            removeConnection(userId);
        };
    }

    private static Consumer<Throwable> onErrorCallback(String userId) {
        return throwable -> {
            System.out.println("==================================================");
            System.out.println("throwable.getMessage() = " + throwable.getMessage());
            log.error("SSE 连接发生异常，用户：{}", userId);
            removeConnection(userId);
        };
    }

    private static Runnable onCompletionCallback (String userId) {
        return  () -> {
            log.info("当前SSE连接完成并结束，用户：{}", userId);
            removeConnection(userId);
        };
    }

    private static void removeConnection(String userId) {
        clientCount.getAndDecrement();
        log.info("SSE 链接被移除，移除用户 ID 为: {}", userId);
        clients.remove(userId);
    }

    public static Integer getClientCount() {
        return clientCount.intValue();
    }

}
